#include "stratified.h"
#include "../core/sampling.h"
#include "../core/point.h"

StratifiedSampelr::StratifiedSampelr(int xstart, int xend, int ystart, int yend, 
									int xs, int ys, bool lns, bool ts, bool jitter) :
					Sampler(xstart, xend, ystart, yend, xs * ys, lns, ts) {
	xPos = xstart;
	yPos = ystart;
	xPixelSamples = xs;
	yPixelSamples = ys;
	jitterSamples = jitter;
	rng = new RNG(xstart + ystart *(xend-xstart));

	// Allocate memory for samples
	imageSamples = new Point2[xPixelSamples * yPixelSamples];
	lensSamples = new Point2[xPixelSamples * yPixelSamples];
	timeSamples = new float[xPixelSamples * yPixelSamples];

	memset(lensSamples, 0, sizeof(Point2)*samplePerPixel);
	memset(timeSamples, 0, sizeof(float)*samplePerPixel);	
}

StratifiedSampelr::~StratifiedSampelr() {
	SAFE_DELETE_ARRAY(imageSamples);
	SAFE_DELETE_ARRAY(lensSamples);
	SAFE_DELETE_ARRAY(timeSamples);
	SAFE_DELETE_ARRAY(rng);
}

Sampler *StratifiedSampelr::GetSubSampler(int num, int count) {
	int x0, x1, y0, y1;
	ComputeSubWindow(num, count, &x0, &x1, &y0, &y1);
	if (x0 == x1 || y0 == y1) 
		return NULL;
	return new StratifiedSampelr(x0, x1, y0, y1, xPixelSamples, yPixelSamples, needLensSamples, needTimeSamples, jitterSamples);
}

int StratifiedSampelr::GetMoreSamples(Sample *sample) {
	if (yPos == yPixelEnd) 
		return 0;
	
	int nSamples = xPixelSamples * yPixelSamples;

	// Generate stratified samples
	StratifiedSample2D(imageSamples, xPixelSamples, yPixelSamples, jitterSamples, rng);
	if(needLensSamples)
		StratifiedSample2D(lensSamples, xPixelSamples, yPixelSamples, jitterSamples, rng);
	if(needTimeSamples)
		StratifiedSample1D(timeSamples, nSamples, jitterSamples, rng);
	
	// Decorrelate sample dimensions
	if(needLensSamples)
		Shuffle(lensSamples, nSamples, rng);
	if(needTimeSamples)
		Shuffle(timeSamples, nSamples, rng);
	
	// Get samples
	for(int i = 0; i < nSamples; i++) {
		sample[i].imageSample = imageSamples[i] + Point2(xPos, yPos);
		sample[i].lensSample = lensSamples[i];
		sample[i].time = timeSamples[i];
		// Generate samples for integrator
		for(int j = 0; j < sample[i].n1D.size(); j++)
			LatinHypercube1D(sample[i].oneD[j], sample[i].n1D[j], rng);
		for(int j = 0; j < sample[i].n2D.size(); j++)
			LatinHypercube2D(sample[i].twoD[j], sample[i].n2D[j], rng);
	}
	
	// advance
	xPos++;
	if(xPos == xPixelEnd) {
		xPos = xPixelStart;
		yPos++;
	}
	
	return nSamples;
}